package com.xiashitech.agent.instrumentation.spring.security;

import com.xiashitech.agent.config.StartConfig;
import com.xiashitech.agent.constants.Instrumentation;
import com.xiashitech.agent.constants.SpecialConfigString;
import com.xiashitech.agent.dto.UserRequest;
import com.xiashitech.agent.instrumentation.base.AbstractAttributesExtractor;
import com.xiashitech.agent.instrumentation.base.context.TraceContextVars;
import com.xiashitech.agent.instrumentation.base.context.bridge.BridgeBaggageBuilder;
import com.xiashitech.agent.instrumentation.bridge.Java8Bridge;
import com.xiashitech.agent.instrumentation.utils.BaggageUtil;
import com.xiashitech.agent.instrumentation.utils.MachineAttributesUtil;
import com.xiashitech.agent.utils.JsonUtil;
import com.xiashitech.agent.utils.SessionUtil;
import com.xiashitech.interfaces.agent.dto.AgentConfigDTO;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.api.baggage.Baggage;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.api.baggage.BaggageEntry;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.api.common.AttributesBuilder;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Context;
import ognl.Ognl;
import org.jetbrains.annotations.Nullable;
import org.springframework.http.ResponseEntity;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.security.oauth2.common.OAuth2AccessToken;
import org.springframework.security.web.server.WebFilterExchange;

import java.util.Iterator;
import java.util.Map;
import java.util.Optional;

public class XSSpringSecurityAttributesExtractor extends AbstractAttributesExtractor<UserRequest,Object> {
    public String getAttributesExtractorName() {
        return XSSpringSecurityAttributesExtractor.class.getName();
    }
    @Override
    public void onXSStart(AttributesBuilder attributesBuilder, Context context, UserRequest userRequest) {
        try {
            if(userRequest != null) {
                if(userRequest.getWebFilterExchange() != null) {
                    String userId = userRequest.getUserid();
                    ServerHttpRequest serverHttpRequest = ((WebFilterExchange)userRequest.getWebFilterExchange()).getExchange().getRequest().mutate().headers(httpHeaders -> {
                        httpHeaders.set(Instrumentation.USERID, userId);
                    }).build();
                    ((WebFilterExchange)userRequest.getWebFilterExchange()).getExchange().mutate().request(serverHttpRequest);
                }
                setAttributes(attributesBuilder,userRequest);
            }
        } catch (Throwable e) {
            StartConfig.throwableStackPrinter(e);
        }
    }

    @Override
    public void onXSEnd(AttributesBuilder attributesBuilder, Context context, UserRequest userRequest, @Nullable Object objReturn, @Nullable Throwable throwable) {
        try {
            if(objReturn != null && objReturn instanceof ResponseEntity) {
                if(throwable == null) {
                    ResponseEntity responseEntity = (ResponseEntity)objReturn;
                    if(responseEntity.hasBody()) {
                        Object responseEntityBody = responseEntity.getBody();
                        if(responseEntityBody instanceof OAuth2AccessToken) {
                            OAuth2AccessToken oAuth2AccessToken = (OAuth2AccessToken)responseEntityBody;
                            Optional<UserRequest> userRequestOptional = Java8Bridge.executeRuleForUserInfoMatchRule(oAuth2AccessToken);
                            if(userRequestOptional.isPresent()) {
                                Boolean sessionUpdate = userRequest.getSessionUpdate();
                                userRequest = userRequestOptional.get();
                                userRequest.setSessionUpdate(sessionUpdate);
                                setAttributes(attributesBuilder,userRequest);
                            }
                        }
                    }
                }
            } else if(objReturn != null && objReturn instanceof UserRequest) {
                UserRequest userRequestForReturn = (UserRequest)objReturn;
                setAttributes(attributesBuilder,userRequestForReturn);
                attributesBuilder.put("xs.loginRequest",userRequestForReturn.getLoginRequest() == null ? false : userRequestForReturn.getLoginRequest().booleanValue());
                attributesBuilder.put("xs.loginName",userRequestForReturn.getLoginName());
                attributesBuilder.put("xs.sessionId",userRequestForReturn.getSessionId());
            } else if (objReturn != null) {
                // 针对 SecurityContext 的 getAuthentication 插桩点
                if(Class.forName("org.springframework.security.core.Authentication").isAssignableFrom(objReturn.getClass())) {
                    String user = getFromAuthentication(objReturn);
                    if(user != null && !user.isEmpty()) { //  && !"anonymousUser".equalsIgnoreCase(user)
                        userRequest.setUserid(user);
                        setAttributes(attributesBuilder, userRequest);
                    }
                }

                // 针对 SecurityContextHolder 的 setContext 插桩点
                if(Class.forName("org.springframework.security.core.context.SecurityContext").isAssignableFrom(objReturn.getClass())) {
                    Object authentication = objReturn.getClass().getMethod("getAuthentication").invoke(objReturn);
                    String user = getFromAuthentication(authentication);
                    if(user != null && !user.isEmpty()) { //  && !"anonymousUser".equalsIgnoreCase(user)
                        userRequest.setUserid(user);
                        setAttributes(attributesBuilder, userRequest);
                    }
                }
            }
        } catch (Throwable e) {
            StartConfig.throwableStackPrinter(e);
        }
        MachineAttributesUtil.fillMachineAttributes(attributesBuilder);
    }

    private void setAttributes(AttributesBuilder attributesBuilder,UserRequest userRequest) {
        attributesBuilder.put("xs.userId",userRequest.getUserid());
        attributesBuilder.put("xs.requestDate",userRequest.getRequestDate());
        attributesBuilder.put("xs.sessionUpdate",userRequest.getSessionUpdate());
        // 不能 makeCurrent，会引发 trace 混合，需要存到 bridgeBaggage
        Baggage baggage = SessionUtil.syncBaggage().toBuilder()
                .put("xs.userId",userRequest.getUserid())
                .put("xs.requestDate",userRequest.getRequestDate())
                .put("xs.sessionUpdate",userRequest.getSessionUpdate() == null ? "false" : userRequest.getSessionUpdate().toString())
                .build();//.makeCurrent();

        TraceContextVars contextVars = TraceContextVars.get();
        Map<String, BaggageEntry> mapBaggage = baggage.asMap();
        Iterator<Map.Entry<String, BaggageEntry>> iteratorBaggage = mapBaggage.entrySet().iterator();
        while (iteratorBaggage.hasNext()) {
            Map.Entry<String, BaggageEntry> entry = iteratorBaggage.next();
            String val = "";
            BaggageEntry be = entry.getValue();
            if(be != null) {
                val = be.getValue();
                if(val == null)
                    val = "";
            }
            contextVars.set(entry.getKey(),val);
        }
    }

    private String getFromAuthentication(Object authentication) {
        String ret = "";
        if(authentication == null)
            return ret;
        try {
            if(!(Class.forName("org.springframework.security.core.Authentication").isAssignableFrom(authentication.getClass()))) {
                return ret;
            }

            Object principal = authentication.getClass().getMethod("getPrincipal").invoke(authentication);
            if(principal == null) {
                return ret;
            }

            // 使用 userIdPath 为特定字符串来决定是否把所有数据取上来
            String userIdPath = AgentConfigDTO.getConfig().getUseridpath().get(StartConfig.getParamValueFromKey(StartConfig.systemId));
            if (SpecialConfigString.ALL_AUTHENTICATION.getConfiString().equalsIgnoreCase(userIdPath)) { // 这个配置，因为会把整个 authentication 放到 http 头信息中，可能会破坏 http 请求
                ret = principal.toString() + "##" + JsonUtil.convertObjectToString(authentication);
            } else {
                if (principal.getClass() == String.class)
                    ret = principal.toString();
                else // SOC 系统：userIdPath 配置为：username，如果取 id，就配 id
                    try {
                        ret = (String) Ognl.getValue(userIdPath, principal);
                    } catch (Exception e) {
                        StartConfig.throwableStackPrinter(e);
                    }
            }
        } catch (Exception ee) {
            StartConfig.throwableStackPrinter(ee);
        }
        return ret;
    }

}
